
#include <cmath>
#include <vector>
#include <memory>
#include <stdlib.h>
#include <iostream>
#include <iterator>
#include <algorithm>
#include <opencv2/opencv.hpp>

using namespace cv;


#include <Eigen/Dense>
#include <Eigen/SparseCore>
#include <Eigen/SparseCholesky>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Sparse>

#include "native_debug.h"


namespace r3d_depth_image {
	class FastBilateralSolverFilter_
	{
	public:

		static Ptr<FastBilateralSolverFilter_> create(
			InputArray guide,
			InputArray depth,
			double sigma_spatial,
			double sigma_luma,
			double sigma_chroma,
			double sigma_depth,
			double lambda,
			int num_iter,
			double max_tol,
			const cv::Mat & process_mask,
			int stride,
			const cv::Mat & edge_mask) {
			if (guide.type() != CV_8UC1 && guide.type() != CV_8UC3) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter::create()]: 1st arg illegal");
				exit(-1);
			}
			FastBilateralSolverFilter_ *fbs = new FastBilateralSolverFilter_();
			cv::Mat gui = guide.getMat();
			cv::Mat depth_map = depth.getMat();
			fbs->init(gui, depth_map, sigma_spatial, sigma_luma, sigma_chroma, sigma_depth, lambda,
				num_iter, max_tol, process_mask, stride, edge_mask);
			return cv::Ptr<FastBilateralSolverFilter_>(fbs);
		}

		void init(cv::Mat& reference,
			cv::Mat& depth,
			double sigma_spatial,
			double sigma_luma,
			double sigma_chroma,
			double simga_depth,
			double lambda,
			int max_iter,
			double max_tol,
			const cv::Mat & process_mask,
			int stride,
			const cv::Mat & edge_mask);

		void filter(InputArray src, InputArray confidence, OutputArray dst) {
			IMAGE_DEPTH_LOGI("FastBilateralSolverFilter_::filter() begin");
			if (src.empty() ||
				(src.depth() != CV_8U &&
					src.depth() != CV_16S &&
					src.depth() != CV_16U &&
					src.depth() != CV_32F) ||
				src.channels() > 4) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter_::filter()]: 1st arg illegal");
				exit(-1);
			}
			if (confidence.empty() ||
				(confidence.depth() != CV_8U &&
					confidence.depth() != CV_32F) ||
				confidence.channels() != 1) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter_::filter()]: 2nd arg illegal");
				exit(-1);
			}
			if (src.rows() != rows || src.cols() != cols) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter:_::filter()]: 1st arg illegal");
				exit(-1);
			}
			if (confidence.rows() != rows || confidence.cols() != cols) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter_::filter()]: 2nd arg illegal");
				exit(-1);
			}

			std::vector<Mat> src_channels;
			std::vector<Mat> dst_channels;
			if (src.channels() == 1)
				src_channels.push_back(src.getMat());
			else
				cv::split(src, src_channels);

			Mat conf = confidence.getMat();

			for (int i = 0; i < src.channels(); i++) {
				Mat cur_res = src_channels[i].clone();
				solve(cur_res, conf, cur_res);
				cur_res.convertTo(cur_res, src.type());
				dst_channels.push_back(cur_res);
			}

			dst.create(src.size(), src_channels[0].type());
			if (src.channels() == 1) {
				Mat& dstMat = dst.getMatRef();
				dstMat = dst_channels[0];
			}
			else merge(dst_channels, dst);
			if (src.type() != dst.type() || src.size() != dst.size()) {
				IMAGE_DEPTH_LOGE("[FastBilateralSolverFilter_::filter()]: Interal Error!");
				exit(-1);
			}
			IMAGE_DEPTH_LOGE("FastBilateralSolverFilter_::filter() end");
		}

		// protected:
		void solve(cv::Mat& src, cv::Mat& confidence, cv::Mat& dst);
		void Splat(Eigen::VectorXf& input, Eigen::VectorXf& dst);
		void Blur(Eigen::VectorXf& input, Eigen::VectorXf& dst);

		void diagonal(Eigen::VectorXf& v, Eigen::SparseMatrix<float>& mat) {
			mat = Eigen::SparseMatrix<float>(v.size(), v.size());
			for (int i = 0; i < int(v.size()); i++) {
				mat.insert(i, i) = v(i);
			}
		}



	private:
		int npixels;
		int nvertices;
		int cols;
		int rows;
		int num_eq; // number of equations(pixels) to solve
		std::vector<int> splat_idx;
		std::vector<int> full_splat_idx;
		std::vector<std::pair<int, int> > blur_idx;
		Eigen::VectorXf m;
		Eigen::VectorXf n;
		Eigen::SparseMatrix<float, Eigen::ColMajor> blurs;
		Eigen::SparseMatrix<float, Eigen::ColMajor> S;
		Eigen::SparseMatrix<float, Eigen::ColMajor> Dn;
		Eigen::SparseMatrix<float, Eigen::ColMajor> Dm;

		struct bs_params
		{
			float lam;
			float cg_tol;
			int cg_maxiter;
			bs_params() {
				lam = 128.0;
				cg_tol = 1e-5f;
				cg_maxiter = 25;
			}
		};
		bs_params bs_param;

	};
};